"""
@brief The websocket client for DC Motor application
"""
import json
import struct
import sys

from umodbus.serial import Serial
from umodbus import functions

from ws_connection import ClientClosedError
from ws_server import WebSocketClient


class MotorModbusMaster(Serial):
    """Rewrite some methods, so it fit into the motor application."""
    MB_M1_BASE = 0x1200
    
    def read_holding_registers(self, 
        slave_addr: int, 
        starting_addr: int, 
        register_qty: int) -> bytes:
        modbus_pdu = functions.read_holding_registers(starting_address=starting_addr, quantity=register_qty)
        
        # 参数 count 表示返回的 response 是否包含 count
        return self._send_receive(modbus_pdu=modbus_pdu, slave_addr=slave_addr, count=True)
    
    @staticmethod
    def parse_motor_pos_spd(buf: bytes) -> tuple[int, ...]:
        """解析电机数据 PosTar, SpdTar, PosFb, SpdFb"""
        # reconstruct buf's word order
        # modbus use big-edian
        _bs = bytearray()
        for i in range(4):
            _bs.extend(buf[i*4+2 : i*4+4])  # high 16-bit
            _bs.extend(buf[i*4 : i*4+2])   # low 16-bitprint(_bs)
        return struct.unpack(">iiii", _bs)

    def mb_poll_motor_info(self):
        _minfo = []
        for i in range(4):
            _resp = self.read_holding_registers(
                slave_addr = 1,
                starting_addr = self.MB_M1_BASE + i*0x200,
                register_qty = 8,
            )
            _minfo.append(self.parse_motor_pos_spd(_resp))
        
        return _minfo
    
    def mb_set_motor(self, motor_num, pos_tar, spd_tar, ctrl_mod):
        print("set:")
        print(motor_num, pos_tar, spd_tar, ctrl_mod)
        # addr of PosTar
        _reg_addr = self.MB_M1_BASE + (motor_num-1) * 0x200 + 0x04
        _regs = (pos_tar&0xFFFF, (pos_tar>>16)&0xFFFF, spd_tar&0xFFFF, (spd_tar>>16)&0xFFFF)
        self.write_multiple_registers(
            slave_addr=1,
            starting_address = _reg_addr,
            register_values = _regs,
        )
        
        #addr of MotCtrlMod
        _reg_addr = self.MB_M1_BASE + (motor_num-1) * 0x200 + 0x09
        _regs = (ctrl_mod,)
        self.write_multiple_registers(
            slave_addr=1,
            starting_address = _reg_addr,
            register_values = _regs,
        )
        

class WsModbusClient(WebSocketClient):
    """与电机Modbus通信"""
    def __init__(self, conn):
        super().__init__(conn)
        self.modbus = MotorModbusMaster(
            pins=(17, 16),  # (TX, RX)
            baudrate = 9600,
            uart_id = 2,  # use UART2
        ) 
    
    def _ws_recv(self, data):
        if data is None:
            return
        print("ws recv:")
        print(data)
        try:
            data = json.loads(data)
            self.modbus.mb_set_motor(
                motor_num = int(data[0]),
                pos_tar = int(data[1]),
                spd_tar = int(data[2]),
                ctrl_mod = int(data[3])
            )
        except ValueError as e:
            print(e)  # JSON error
    
    def process(self):
        # ===== Send received data to UART0 =====
        try:
            _recv = self.connection.read()
        except ClientClosedError:
            self.connection.close()
            return
        
        self._ws_recv(_recv)
        
        try:
            _data = self.modbus.mb_poll_motor_info()
            self.connection.write(str(_data))
            print(_data)
        except ValueError as e:
            print(e)
        except OSError as e:
            print(e)
        



